Allow position_ids_start=2 on DataCollatorWithFlattening for RoBERTa etc. - #47525
Allow position_ids_start=2 on DataCollatorWithFlattening for RoBERTa etc.#47525tomaarsen wants to merge 3 commits into
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| indices_q = (position_ids == 0).nonzero().view(-1) | ||
| # Packed sequences all restart from the same first position id, but it is not always 0 | ||
| # (RoBERTa-like models start at padding_idx + 1) | ||
| indices_q = (position_ids == position_ids.min()).nonzero().view(-1) |
There was a problem hiding this comment.
Can we add an integration test for roberta that tests this case implicitly? It's ok to require FA or kernels
|
[For maintainers] Suggested jobs to run (before merge) run-slow: roberta |
CI recapDashboard: View test results in Grafana |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
In general yes, just some comments for the test design
| def test_flash_attn_padding_free_position_ids_start(self): | ||
| """Test that flattening a batch with position_ids_start=pad_token_id + 1 matches the padded batched run.""" | ||
| config = self.model_tester.prepare_config_and_inputs()[0] | ||
| config.attn_implementation = "flash_attention_2" |
There was a problem hiding this comment.
let's use set_attn_implementation method instead of setting in the config
| # Sample from [2, vocab_size) so no token collides with pad_token_id, which would shift | ||
| # the position ids the model computes internally in the batched reference run | ||
| seq_len_a, seq_len_b = 3, 4 | ||
| input_ids = ids_tensor([2, seq_len_b], config.vocab_size - 2) + 2 | ||
| attention_mask = torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]], device=torch_device) | ||
| input_ids[attention_mask == 0] = config.pad_token_id |
There was a problem hiding this comment.
Why not use the prepared inputs and just use a modified attention mask?
| # bf16 noise between the padded and packed kernel shapes reaches ~3e-2, wrong position ids differ by >1.9 | ||
| torch.testing.assert_close(batched[0, :seq_len_a], flattened[0, :seq_len_a], atol=1e-1, rtol=1e-1) | ||
| torch.testing.assert_close(batched[1, :seq_len_b], flattened[0, seq_len_a:], atol=1e-1, rtol=1e-1) |
There was a problem hiding this comment.
I'd rather use fp16 then + a lower tol, this seems quite high no?
3 participants
Fixes #47496
What does this PR do?
Pull Request overview
position_ids_startargument toDataCollatorWithFlatteningposition_ids == position_ids.min()instead ofposition_ids == 0Details
As described in #47496,
DataCollatorWithFlatteningnumbers positions from 0, but RoBERTa-like architectures (roberta,xlm_roberta,camembert,mpnet,esmwith absolute position embeddings, and more) compute their position ids starting atpadding_idx + 1(usually 2) when none are passed. Feeding these models the collator's 0-basedposition_idssilently shifts the position embeddings and degrades the outputs. This affects some very common checkpoints, e.g.xlm-roberta.This PR adds a
position_ids_startargument (default 0, existing behavior is unchanged) so users can match the numbering of their model, e.g.position_ids_start=2for XLM-R. The collator cannot reliably derive this from a config:mpnetfor example hardcodespadding_idx = 1regardless ofconfig.pad_token_id. The docstring now documents all arguments and warns about the pitfall.The new argument alone turned out to be insufficient on the path where FlashAttention infers the varlen kwargs from
position_idsat runtime (i.e. withoutreturn_flash_attn_kwargs=True, like the reproduction in the issue).prepare_fa_kwargs_from_position_idsdetected sequence starts viaposition_ids == 0, which never matches offset positions, after which it crashes on an emptycu_seq_lens.diff().max(). I changed the detection toposition_ids == position_ids.min(), mirroring_is_packed_sequencewhich already handles a non-zero start viaposition_ids.min(). Whenever a 0 position is present, this is behavior-identical, so existing models are unaffected.I verified that with
position_ids_start=2the collator output exactly matchesRobertaEmbeddings.create_position_ids_from_input_idson unpadded sequences, and that the RoBERTa embedding layer produces identical outputs for the flattened batch and for the individual sequences. The flattening collator tests now cover the new argument as well as the boundary inference on collator output.Because of this
prepare_fa_kwargs_from_position_idspatch, the PR got a bit bigger than I originally intended in #47496 (comment), but it does seem important.Code Agent Policy
The tests and docstrings were written by Fable and reviewed by myself.
Before submitting
Pull Request checks?
to it if that's the case.
Who can review?
@vasqu